import threading
import queue
import time

class ThreadPool:
    def __init__(self, num_threads):
        # 创建任务队列
        self.task_queue = queue.Queue()
        # 创建线程池
        self.threads = []
        # 创建并启动工作线程
        for _ in range(num_threads):
            thread = threading.Thread(target=self._worker)
            thread.daemon = True  # 设置为守护线程，主程序退出时自动结束
            thread.start()
            self.threads.append(thread)
    
    def _worker(self):
        """工作线程的执行函数"""
        while True:
            # 从队列获取任务 (阻塞式获取)
            task = self.task_queue.get()
            if task is None:  # 收到终止信号
                break
            func, args, kwargs = task
            try:
                # 执行任务
                func(*args, **kwargs)
            except Exception as e:
                print(f"Task failed with error: {e}")
            finally:
                # 通知队列任务已完成
                self.task_queue.task_done()
    
    def add_task(self, func, *args, **kwargs):
        """添加任务到队列"""
        self.task_queue.put((func, args, kwargs))
    
    def wait_completion(self):
        """等待所有任务完成"""
        self.task_queue.join()
    
    def shutdown(self):
        """关闭线程池（等待所有任务完成后退出）"""
        # 等待剩余任务完成
        self.wait_completion()
        # 发送终止信号给所有线程
        for _ in range(len(self.threads)):
            self.task_queue.put(None)
        # 等待所有线程结束
        for thread in self.threads:
            thread.join()

# 测试函数
def sample_task(task_id, delay=1):
    print(f"Task {task_id} started")
    time.sleep(delay)
    print(f"Task {task_id} completed")

if __name__ == "__main__":
    # 创建包含3个工作线程的线程池
    pool = ThreadPool(3)
    
    # 添加10个任务
    for i in range(10):
        pool.add_task(sample_task, i, delay=1)
    
    # 等待所有任务完成
    pool.wait_completion()
    print("\nAll tasks completed!")
    
    # 添加更多任务
    for i in range(10, 15):
        pool.add_task(sample_task, i, delay=0.5)
    
    # 关闭线程池（会再次等待任务完成）
    pool.shutdown()
    print("Thread pool shutdown")